# uncompyle6 version 3.2.3
# Python bytecode 3.6 (3379)
# Decompiled from: Python 3.6.8 |Anaconda custom (64-bit)| (default, Feb 21 2019, 18:30:04) [MSC v.1916 64 bit (AMD64)]
# Embedded file name: helpers\scan_result_handling\__init__.py
import logging
from sqlalchemy import select as sql_select
from sqlalchemy import func as sql_func
import sqlalchemy.exc
from db.tables.vuln_types import VulnTypeRow
from db.tables.scan_sessions import ScanSessionRow, ScanSessionsTable
from db.tables.scan_session_jobs import ScanSessionJobRow
from db.tables.scans import ScanRow, ScansTable
from db.tables.targets import TargetsTable, TargetRow
from db.tables.target_vulns import TargetVulnsTable, TargetVulnRow
from helpers.address import extract_domain_from_address, extract_host_from_address

logger = logging.getLogger("helpers.scan_result_handling")
scan_sessions_cache = dict()
vulnerability_type_cache = dict()


def get_target_id_for_domain(scan_session_id, domain, connection=None):
    """
    attempt to retrieve the target_id for a certain scan_session_id and domain combination,
    if connection is none will retrieve data from cache only
    :param scan_session_id:
    :param domain:
    :param connection:
    :return:
    """
    scan_session = get_scan_data(scan_session_id, connection)
    if not scan_session:
        return
    else:
        if domain.lower() == scan_session["domain"].lower():
            target_id = scan_session["target_id"]
        else:
            target_id = scan_session["allowed_targets"].get(domain)
        return target_id


def get_scan_data(scan_session_id, connection=None):
    """
    
    :param scan_session_id:
    :param connection:
    :return:
    """
    result = scan_sessions_cache.get(scan_session_id)
    if result is None:
        if connection is None:
            return result
        select_from = ScanSessionsTable.join(
            ScansTable, ScanRow.scan_id == ScanSessionRow.scan_id
        )
        select_from = select_from.join(
            TargetsTable, TargetRow.target_id == ScanRow.target_id
        )
        q = (
            sql_select(
                (
                    ScanSessionRow,
                    ScanRow.owner_id,
                    TargetRow.criticality,
                    TargetRow.address,
                    ScanRow.owner_id,
                    ScanRow.continuous,
                    ScanRow.creator_id,
                )
            )
            .select_from(select_from)
            .where(ScanSessionRow.scan_session_id == scan_session_id)
        )
        scan_session = connection.execute(q).fetchone()
        if not scan_session:
            result = False
        else:
            result = dict(scan_session)
            if result["allowed_targets"] is None:
                result["allowed_targets"] = dict()
            result["domain"] = extract_domain_from_address(result["address"])
            result["host"] = extract_host_from_address(result["address"])
        scan_sessions_cache[scan_session_id] = result
    return result


def get_vulnerability_type(vt_id, connection=None):
    result = vulnerability_type_cache.get(vt_id)
    if not result:
        if connection is None:
            return result
        result = connection.execute(
            sql_select((VulnTypeRow.name, VulnTypeRow.severity)).where(
                VulnTypeRow.vt_id == vt_id
            )
        ).fetchone()
        if result:
            result = dict(result)
            vulnerability_type_cache[vt_id] = result
    return result


def insert_target_vuln(connection, scan_session_id, target_vuln):
    """
    Attempts to insert a target vuln entry, on collision will update state on vuln status logic
    
    :param connection:
    :param scan_session_id:
    :param target_vuln:
    :return: (vuln_id, rediscovered)
    """
    try:
        query = (TargetVulnsTable.insert().values(**target_vuln)).returning(
            TargetVulnRow.vuln_id
        )
        vuln_id = connection.execute(query).scalar()
        return (vuln_id, False)
    except sqlalchemy.exc.IntegrityError as e:
        str_e = str(e)
        if "target_vulns_target_id_fkey" in str_e:
            return (None, None)
        if "ix_target_vulns_hash" in str_e:
            q = (
                sql_select(
                    (
                        TargetVulnRow.status,
                        TargetVulnRow.criticality,
                        TargetVulnRow.vt_id,
                        TargetVulnRow.vuln_id,
                        TargetVulnRow.rediscovered,
                        TargetVulnRow.scan_session_id,
                    )
                )
                .where(TargetVulnRow.vuln_hash == target_vuln["vuln_hash"])
                .where(TargetVulnRow.target_id == target_vuln["target_id"])
            )
            old_target_vuln = connection.execute(q).fetchone()
            if old_target_vuln:
                rediscovered = False
                if old_target_vuln.status in ("ignored", "false_positive"):
                    logger.info(
                        "skipping ignored vuln, %s:%s",
                        scan_session_id,
                        target_vuln["vuln_hash"],
                    )
                    return (None, None)
                if old_target_vuln.status == "fixed":
                    rediscovered = True
                if old_target_vuln.scan_session_id == scan_session_id:
                    logger.debug("vuln hash collision")
                    if old_target_vuln.status == "open":
                        return (old_target_vuln.vuln_id, rediscovered)
                    criticality = (
                        old_target_vuln.criticality
                        if old_target_vuln.criticality
                        else 0
                    )
                    criticality = max(criticality, target_vuln["criticality"])
                    update_values = dict(
                        scan_session_id=scan_session_id,
                        last_seen=sql_func.current_timestamp(),
                        request=target_vuln["request"],
                        status="open",
                        continuous=target_vuln["continuous"],
                        criticality=criticality,
                        sensor_details=target_vuln.get("sensor_details"),
                    )
                    if not old_target_vuln.rediscovered:
                        if rediscovered:
                            update_values["rediscovered"] = rediscovered
                    q = (
                        (TargetVulnsTable.update(values=update_values))
                        .where(TargetVulnRow.target_id == target_vuln["target_id"])
                        .where(TargetVulnRow.vuln_hash == target_vuln["vuln_hash"])
                    )
                    connection.execute(q)
                    return (old_target_vuln.vuln_id, rediscovered)
            else:
                raise
